import os
import time
import numpy as np
import matplotlib.pyplot as plt

SAVE_DIR = os.getcwd()
RATIO = 0.2
np.random.seed((int(time.time())))

def buffon_needle(a=4, l=3, n=1e6):
    n = int(n)
    phi_array = np.pi * np.random.random(n)
    x_array = 0.5 * a * np.random.random(n)
    status_array = 0.5 * l * np.sin(phi_array) - x_array
    status_array = np.where(status_array > 0, 1, 0)
    
    status_sum = 0
    p_array = []
    n_array = np.linspace(1, n, n, dtype=np.int64)
    for _ in n_array:
        status_sum += status_array[_ - 1]
        p_array.append(status_sum / _)
    p_array = np.array(p_array)
    p_array = np.where(p_array == 0, np.inf, p_array)
    pi_array = 2 * l / (p_array * a)
    delta_pi_array = 2 * np.pi * np.sqrt((1 - p_array) / (p_array * n_array))
    epsilon_array = 2 * np.sqrt((p_array - np.square(p_array)) / n_array)
    
    plt.figure(figsize=(16, 9))
    plt.plot(n_array, pi_array, label="${\pi}_{detect}$")
    plt.plot([n_array[0], n_array[-1]], [np.pi] * 2, label="${\pi}_{real}$")
    plt.ylim((1 - RATIO) * np.pi, (1 + RATIO) * np.pi)
    plt.xlabel("The Number of Throws $N$")
    plt.ylabel("The Value of $\pi$")
    plt.title("$\pi \sim N$")
    plt.legend()
    plt.savefig(os.path.join(SAVE_DIR, "value.png"))
    
    plt.cla()
    plt.plot(n_array, delta_pi_array, label="${\Delta \pi}_{theorectical}$")
    plt.plot(n_array, epsilon_array, label="${\Delta \pi}_{real}$")
    plt.xlabel("The Number of Throws $N$")
    plt.ylabel("The Value of $\Delta \pi$")
    plt.ylim(-0.05, 0.25)
    plt.title("$\Delta \pi \sim N$")
    plt.legend()
    plt.savefig(os.path.join(SAVE_DIR, "error.png"))
    
if __name__ == "__main__":
    buffon_needle()